﻿using FlyRemote.Base;
using Lyo.Comm;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;

namespace FlyRemote.Remote
{

    /// <summary>
    /// 通过 获取 目标程序的 peb 数据 , 来遍历模块  , 或 遍历栈  
    /// </summary>
    public static class PebHelper
    {



        #region  64位数据 


        public struct UNICODE_STRING64
        {
            public ushort Length;
            public ushort MaximumLength;
            public ulong Buffer;
        }

        public struct LDR_DATA_TABLE_ENTRY64
        {
            public LIST_ENTRY64 InLoadOrderLinks;
            public LIST_ENTRY64 InMemoryOrderModuleList;
            public LIST_ENTRY64 InInitializationOrderModuleList;
            public ulong DllBase;
            public ulong EntryPoint;
            public ulong SizeOfImage;
            public UNICODE_STRING64 FullDllName;
            public UNICODE_STRING64 BaseDllName;

        }

        public struct LIST_ENTRY64
        {
            public ulong Flink;
            public ulong Blink;
        }

        public struct PEB_LDR_DATA64
        {
            public uint Length;
            public byte Initialized;
            public IntPtr SsHandle;
            public LIST_ENTRY64 InLoadOrderModuleList;
        }


        public struct PEB64
        {
            public byte InheritedAddressSpace;
            public byte ReadImageFileExecOptions;
            public byte BeingDebugged;
            public byte BitField;
            public IntPtr Mutant;
            public IntPtr ImageBaseAddress;
            public ulong Ldr;
        }

        public struct PROCESS_BASIC_INFORMATION64
        {
            public ulong ExitStatus;
            public ulong PebBaseAddress;
            public ulong AffinityMask;
            public ulong BasePriority;
            public ulong UniqueProcessId;
            public ulong InheritedFromUniqueProcessId;
        }






        ///  这里是线程相关  



        public struct NT_TIB64
        {
            public ulong ExceptionList;
            public ulong StackBase;
            public ulong StackLimit;
            public ulong SubSystemTib;

        }

        public struct THREAD_BASIC_INFORMATION64
        {
            public ulong ExitStatus;
            public ulong TebBaseAddress;
            public ulong ClientId;
            public ulong AffinityMask;
            public ulong Priority;
            public ulong BasePriority;
        }


        #endregion


        #region 32位数据 -----------------------------

        public struct UNICODE_STRING32
        {
            public ushort Length;
            public ushort MaximumLength;
            public uint Buffer;
        }

        public struct LDR_DATA_TABLE_ENTRY32
        {
            public LIST_ENTRY32 InLoadOrderLinks;
            public LIST_ENTRY32 InMemoryOrderModuleList;
            public LIST_ENTRY32 InInitializationOrderModuleList;
            public uint DllBase;
            public uint EntryPoint;
            public uint SizeOfImage;
            public UNICODE_STRING32 FullDllName;
            public UNICODE_STRING32 BaseDllName;
            public uint Flags;
            public ushort LoadCount;
            public ushort TlsIndex;

            /* union */
            public LIST_ENTRY32 HashLinks;
            public uint SectionPointer;
            public uint CheckSum;
            //union
            public uint TimeDateStamp;
            public uint LoadedImports;
            public uint EntryPointActivationContext;
            public uint PatchInformation;

        }

        public struct LIST_ENTRY32
        {
            public uint Flink;
            public uint Blink;
        }

        public struct PEB_LDR_DATA32
        {
            public uint Length;
            public byte Initialized;
            public uint SsHandle;
            public LIST_ENTRY32 InLoadOrderModuleList;
        }


        public struct PEB32
        {
            public byte InheritedAddressSpace;
            public byte ReadImageFileExecOptions;
            public byte BeingDebugged;
            public byte BitField;
            public uint Mutant;
            public uint ImageBaseAddress;
            public uint Ldr;
        }
        public struct PROCESS_BASIC_INFORMATION32
        {
            public uint ExitStatus;
            public uint PebBaseAddress;
            public uint AffinityMask;
            public uint BasePriority;
            public uint UniqueProcessId;
            public uint InheritedFromUniqueProcessId;
        }


        public struct NT_TIB32
        {
            public uint ExceptionList;
            public uint StackBase;
            public uint StackLimit;
            public uint SubSystemTib;

        }


        #endregion



        // #pragma warning restore 169, 414



        [DllImport("ntdll.dll", SetLastError = true)]
        static extern int NtQueryInformationProcess(IntPtr Handle, int Class, IntPtr data, int Length, IntPtr returnLength);




        [DllImport("ntdll.dll", SetLastError = true)]
        static extern int NtQueryInformationThread(IntPtr Handle, int Class, IntPtr data, int Length, IntPtr returnLength);


        /// <summary>
        /// 通过查询 peb , 获取目标程序的 dll 加载情况 
        /// </summary>
        public static IEnumerable<NormalizedModule> GetModules(IMemoryAssist assist)
        {


            /// !!! 这里要求自身进程 必须是 64位, 如果是32位进程 , 则返回 c0000004 
            if (!Environment.Is64BitProcess)
                throw new Exception("必须是64进程来调用");


            var modules = new List<NormalizedModule>(0x100);

            if (PAFactory.Default.OpenedProcess == null)
                return modules;


            IntPtr data = Marshal.AllocHGlobal(0x1000);        //定义指针
            IntPtr handle = PAFactory.Default.OpenedProcess.Handle;
            bool is32 = PAFactory.Default.IsOpenedProcess32Bit();

            PROCESS_BASIC_INFORMATION64 pbi64 = new PROCESS_BASIC_INFORMATION64();
            var queryStatus = NtQueryInformationProcess(handle, 0, data, Marshal.SizeOf(pbi64), IntPtr.Zero);

            if (queryStatus != 0)
            {

                Marshal.FreeHGlobal(data);
                return modules;
            }
            pbi64 = StructHelper.FromPtr<PROCESS_BASIC_INFORMATION64>(data);

            if (is32)
            {
                /// 如果是 读取32位 , 会有个 0x1000 的偏移 
                ///  其实是分别对应有 32位 64位 两套数据 ，如果是32位， 则64位数据(只有程序初始加载的信息，不全)

                byte[] arr = assist.ReadBytes(pbi64.PebBaseAddress + 0x1000, 0x100, out bool ok);

                PEB32 peb = StructHelper.FromBytes<PEB32>(arr);

                arr = assist.ReadBytes(peb.Ldr, 0x100, out ok);
                PEB_LDR_DATA32 dATA64 = StructHelper.FromBytes<PEB_LDR_DATA32>(arr);

                ulong addr = (ulong)dATA64.InLoadOrderModuleList.Flink;
                LDR_DATA_TABLE_ENTRY32 entry;

                for (; ; addr = entry.InLoadOrderLinks.Flink)
                {
                    arr = assist.ReadBytes(addr, 0x100, out ok);
                    entry = StructHelper.FromBytes<LDR_DATA_TABLE_ENTRY32>(arr);
                    arr = assist.ReadBytes(entry.FullDllName.Buffer, entry.FullDllName.Length, out ok);

                    string dllname = Encoding.Unicode.GetString(arr);

                    // Convert to a normalized module and add it to our list
                    NormalizedModule module = new NormalizedModule(dllname, entry.DllBase, (Int32)entry.SizeOfImage);
                    modules.Add(module);

                    /// 这里已经是最后一个了 
                    if (addr == dATA64.InLoadOrderModuleList.Blink)
                        break;
                }
            }
            else
            {

                byte[] arr = assist.ReadBytes(pbi64.PebBaseAddress, 0x100, out bool ok);

                PEB64 peb = StructHelper.FromBytes<PEB64>(arr);

                arr = assist.ReadBytes(peb.Ldr, 0x100, out ok);
                PEB_LDR_DATA64 dATA64 = StructHelper.FromBytes<PEB_LDR_DATA64>(arr);


                ulong addr = (ulong)dATA64.InLoadOrderModuleList.Flink;
                LDR_DATA_TABLE_ENTRY64 entry = new LDR_DATA_TABLE_ENTRY64();

                for (; ; addr = entry.InLoadOrderLinks.Flink)
                {
                    arr = assist.ReadBytes(addr, 0x100, out ok);
                    entry = StructHelper.FromBytes<LDR_DATA_TABLE_ENTRY64>(arr);
                    arr = assist.ReadBytes(entry.FullDllName.Buffer, entry.FullDllName.Length, out ok);
                    string dllpath = Encoding.Unicode.GetString(arr);


                    // Convert to a normalized module and add it to our list
                    NormalizedModule module = new NormalizedModule(dllpath, entry.DllBase, (Int32)entry.SizeOfImage);
                    modules.Add(module);

                    /// 这里已经是最后一个了 
                    if (addr == dATA64.InLoadOrderModuleList.Blink)
                        break;
                }



            }


            Marshal.FreeHGlobal(data);
            return modules;

        }

        [Flags]
        public enum ThreadAccess : int
        {
            TERMINATE = (0x0001),
            SUSPEND_RESUME = (0x0002),
            GET_CONTEXT = (0x0008),
            SET_CONTEXT = (0x0010),
            SET_INFORMATION = (0x0020),
            QUERY_INFORMATION = (0x0040),
            SET_THREAD_TOKEN = (0x0080),
            IMPERSONATE = (0x0100),
            DIRECT_IMPERSONATION = (0x0200)
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr OpenThread(ThreadAccess dwDesiredAccess, bool bInheritHandle, uint dwThreadId);

        // 栈的地址其实并没有规律, 
        // 目前通过 peb 没有办法遍历 这个 teb 结构 , 
        // 需要分析 x64dbg 实现  -> 使用的是 dbg api 

        // 考虑 NtQueryInformationThread  

        /// <summary>
        /// 通过查询线程信息 , 获取线程当前的栈信息 
        /// </summary>
        public static IEnumerable<NormalizedRegion> GetStacks(IMemoryAssist assist)
        {

            IList<NormalizedRegion> regions = new List<NormalizedRegion>(0x100);

            if (assist == null)
                assist = RMFactory.Default;

            if (PAFactory.Default.OpenedProcess == null)
                return regions;


            /// 需要获取 实时的线程 数据 
            Process p = Process.GetProcessById(PAFactory.Default.OpenedProcess.Id);


            THREAD_BASIC_INFORMATION64 tb64 = new THREAD_BASIC_INFORMATION64();

            bool is32 = PAFactory.Default.IsOpenedProcess32Bit();


            IntPtr data = Marshal.AllocHGlobal(0x1000);



            foreach (ProcessThread thread in p.Threads)
            {
                IntPtr threadHandle = OpenThread(ThreadAccess.QUERY_INFORMATION, false, (uint)thread.Id);

                int status = NtQueryInformationThread(threadHandle, 0, data, Marshal.SizeOf(tb64), IntPtr.Zero);

                if (status == 0)
                {
                    tb64 = StructHelper.FromPtr<THREAD_BASIC_INFORMATION64>(data);

                    if (is32)
                    {
                        uint addr = assist.Read<uint>(tb64.TebBaseAddress, out bool ok);
                        byte[] arr = assist.ReadBytes(addr, 0x100, out ok);

                        NT_TIB32 tib = StructHelper.FromBytes<NT_TIB32>(arr);
                        regions.Add(new NormalizedRegion(tib.StackLimit, (int)(tib.StackBase - tib.StackLimit)));
                    }
                    else
                    {
                        byte[] arr = assist.ReadBytes(tb64.TebBaseAddress, 0x100, out bool ok);
                        NT_TIB64 tib = StructHelper.FromBytes<NT_TIB64>(arr);
                        regions.Add(new NormalizedRegion(tib.StackLimit, (int)(tib.StackBase - tib.StackLimit)));
                    }

                }

            }


            Marshal.FreeHGlobal(data);

            return regions;

        }

    }
}
